# tests/integration/test_conservation.py
"""
Integration Tests for Conservation Laws (守恒律集成测试)
验证守恒归零机制：从任意初态经演化回归 ±0
"""

import pytest
import numpy as np
from typing import Callable, Any
from dataclasses import dataclass

from core.axioms import (
    InnateBalance,
    DualPair,
    ConservationLaw,
    TernaryStructure,
    AxiomSystem
)


# ============================================================================
# 测试用物理系统类
# ============================================================================

@dataclass
class PhysicalState:
    """物理系统状态"""
    energy: float
    momentum: np.ndarray
    position: np.ndarray
    time: float = 0.0
    
    def total_energy(self) -> float:
        """总能量（动能+势能）"""
        kinetic = 0.5 * np.sum(self.momentum ** 2)
        potential = self.energy
        return kinetic + potential
    
    def evolve(self, dt: float, force: Callable[[np.ndarray], np.ndarray]) -> 'PhysicalState':
        """时间演化"""
        # 简单的欧拉法
        new_momentum = self.momentum + force(self.position) * dt
        new_position = self.position + new_momentum * dt
        return PhysicalState(
            energy=self.energy,
            momentum=new_momentum,
            position=new_position,
            time=self.time + dt
        )


@dataclass
class YinYangPair:
    """
    阴阳对：具有守恒律的对偶系统
    阳(+)与阴(-)相互转化，总和守恒归零
    """
    yang: complex  # 阳（正）
    yin: complex   # 阴（负）
    
    def total(self) -> complex:
        """总量：应守恒"""
        return self.yang + self.yin
    
    def evolve(self, rate: float = 0.1) -> 'YinYangPair':
        """
        演化：阳消阴长，阴消阳长
        守恒条件：d(yang)/dt + d(yin)/dt = 0
        """
        # 振荡演化：yang = A*cos(t), yin = -A*cos(t)
        theta = rate
        rotation = np.array([
            [np.cos(theta), -np.sin(theta)],
            [np.sin(theta), np.cos(theta)]
        ])
        
        state = np.array([self.yang.real, self.yin.real])
        new_state = rotation @ state
        
        return YinYangPair(
            yang=complex(new_state[0], self.yang.imag),
            yin=complex(new_state[1], self.yin.imag)
        )
    
    def to_innate_balance(self, tolerance: float = 1e-6) -> bool:
        """检查是否归零到 ±0"""
        total = abs(self.total())
        return total < tolerance


class HarmonicOscillator:
    """谐振子：经典守恒系统"""
    
    def __init__(self, mass: float = 1.0, k: float = 1.0):
        self.mass = mass
        self.k = k  # 弹簧常数
        
    def total_energy(self, position: float, momentum: float) -> float:
        """总能量 E = p²/2m + kx²/2"""
        kinetic = momentum ** 2 / (2 * self.mass)
        potential = 0.5 * self.k * position ** 2
        return kinetic + potential
    
    def evolve(self, state: tuple[float, float], dt: float) -> tuple[float, float]:
        """Hamilton 演化"""
        x, p = state
        # dx/dt = p/m, dp/dt = -kx
        new_x = x + (p / self.mass) * dt
        new_p = p - (self.k * x) * dt
        return new_x, new_p


class QuantumTwoLevel:
    """两能级量子系统"""
    
    def __init__(self, energy_gap: float = 1.0):
        self.E_gap = energy_gap
        
    def total_energy(self, state: np.ndarray) -> float:
        """
        期望能量 <E> = |c₀|² E₀ + |c₁|² E₁
        守恒：|c₀|² + |c₁|² = 1
        """
        c0, c1 = state
        E0, E1 = 0.0, self.E_gap
        return abs(c0) ** 2 * E0 + abs(c1) ** 2 * E1
    
    def evolve(self, state: np.ndarray, dt: float) -> np.ndarray:
        """幺正演化 U = exp(-iHt/ℏ)"""
        c0, c1 = state
        E0, E1 = 0.0, self.E_gap
        
        # 幺正演化矩阵
        phase0 = np.exp(-1j * E0 * dt)
        phase1 = np.exp(-1j * E1 * dt)
        
        new_state = np.array([c0 * phase0, c1 * phase1])
        
        # 归一化（保证概率守恒）
        norm = np.sqrt(np.sum(np.abs(new_state) ** 2))
        return new_state / norm


# ============================================================================
# 测试：基础守恒验证
# ============================================================================

class TestBasicConservation:
    """基础守恒律测试"""
    
    def test_innate_balance_conservation(self):
        """测试先天均衡态的守恒性"""
        balance = InnateBalance()
        
        # 验证总和为零
        total = balance.positive + balance.negative + balance.innate_zero
        assert abs(total) < 1e-10, "先天均衡态总和应为零"
        
        # 验证对称性
        assert balance.verify(), "先天均衡态验证失败"
    
    def test_dual_pair_annihilation(self):
        """测试对偶对湮灭归零"""
        pair = DualPair(positive=1.0, negative=-1.0)
        
        # 湮灭应回归 ±0
        result = pair.annihilate()
        assert isinstance(result, InnateBalance)
        assert result.verify()
    
    def test_energy_conservation_law(self):
        """测试能量守恒律定义"""
        def energy_func(state: PhysicalState) -> float:
            return state.total_energy()
        
        conservation = ConservationLaw(quantity=energy_func, name="Energy")
        
        # 创建初态和末态（能量相同）
        state1 = PhysicalState(
            energy=10.0,
            momentum=np.array([1.0, 0.0]),
            position=np.array([0.0, 0.0])
        )
        state2 = PhysicalState(
            energy=10.0,
            momentum=np.array([0.0, 1.0]),
            position=np.array([1.0, 0.0])
        )
        
        assert conservation.verify(state1, state2, tolerance=1e-6)


# ============================================================================
# 测试：阴阳对守恒演化
# ============================================================================

class TestYinYangConservation:
    """阴阳对守恒机制测试"""
    
    def test_yinyang_total_conservation(self):
        """测试阴阳总量守恒"""
        pair = YinYangPair(yang=1.0+0j, yin=-1.0+0j)
        initial_total = pair.total()
        
        # 经过多步演化
        for _ in range(100):
            pair = pair.evolve(rate=0.1)
        
        final_total = pair.total()
        
        # 总量应保持不变（归零）
        assert abs(initial_total - final_total) < 1e-6, \
            f"阴阳总量不守恒: {initial_total} → {final_total}"
    
    def test_yinyang_periodic_return(self):
        """测试阴阳周期性回归"""
        pair = YinYangPair(yang=1.0+0j, yin=-1.0+0j)
        initial_yang = pair.yang
        
        # 完整周期 (2π)
        steps = int(2 * np.pi / 0.1)
        for _ in range(steps):
            pair = pair.evolve(rate=0.1)
        
        # 应回到初始状态（守恒回归）
        assert abs(pair.yang - initial_yang) < 0.1, \
            "阴阳未能周期回归"
    
    def test_yinyang_zero_convergence(self):
        """测试阴阳归零收敛"""
        # 特殊情况：初始即平衡
        pair = YinYangPair(yang=1.0+0j, yin=-1.0+0j)
        
        assert pair.to_innate_balance(tolerance=1e-6), \
            "平衡态阴阳对应归零于 ±0"


# ============================================================================
# 测试：经典系统守恒
# ============================================================================

class TestClassicalConservation:
    """经典物理系统守恒测试"""
    
    def test_harmonic_oscillator_energy(self):
        """测试谐振子能量守恒"""
        osc = HarmonicOscillator(mass=1.0, k=1.0)
        
        # 初态
        state = (1.0, 0.0)  # (位置, 动量)
        initial_energy = osc.total_energy(*state)
        
        # 演化100步
        dt = 0.01
        for _ in range(100):
            state = osc.evolve(state, dt)
        
        final_energy = osc.total_energy(*state)
        
        # 能量守恒（数值误差在5%内）
        relative_error = abs(final_energy - initial_energy) / initial_energy
        assert relative_error < 0.05, \
            f"谐振子能量不守恒: {initial_energy} → {final_energy}"
    
    def test_harmonic_oscillator_return_to_zero(self):
        """测试谐振子归零机制"""
        osc = HarmonicOscillator(mass=1.0, k=1.0)
        
        # 从非零初态开始
        state = (1.0, 1.0)
        
        # 演化一个完整周期 T = 2π√(m/k)
        T = 2 * np.pi * np.sqrt(osc.mass / osc.k)
        steps = int(T / 0.01)
        dt = 0.01
        
        for _ in range(steps):
            state = osc.evolve(state, dt)
        
        # 应回到原点附近（周期性）
        x, p = state
        assert abs(x - 1.0) < 0.2, "位置未能周期回归"


# ============================================================================
# 测试：量子系统守恒
# ============================================================================

class TestQuantumConservation:
    """量子系统守恒测试"""
    
    def test_two_level_probability_conservation(self):
        """测试两能级系统概率守恒"""
        system = QuantumTwoLevel(energy_gap=1.0)
        
        # 初态：叠加态
        state = np.array([1/np.sqrt(2), 1/np.sqrt(2)], dtype=complex)
        initial_prob = np.sum(np.abs(state) ** 2)
        
        # 演化
        for _ in range(100):
            state = system.evolve(state, dt=0.01)
        
        final_prob = np.sum(np.abs(state) ** 2)
        
        # 概率守恒
        assert abs(final_prob - initial_prob) < 1e-10, \
            "量子概率不守恒"
        assert abs(final_prob - 1.0) < 1e-10, \
            "概率归一化失败"
    
    def test_two_level_energy_expectation(self):
        """测试能量期望值守恒（封闭系统）"""
        system = QuantumTwoLevel(energy_gap=1.0)
        
        # 能量本征态：不随时间变化（除了相位）
        state = np.array([1.0, 0.0], dtype=complex)  # 基态
        initial_energy = system.total_energy(state)
        
        # 演化
        for _ in range(100):
            state = system.evolve(state, dt=0.01)
        
        final_energy = system.total_energy(state)
        
        # 能量期望值守恒
        assert abs(final_energy - initial_energy) < 1e-10, \
            f"量子能量不守恒: {initial_energy} → {final_energy}"


# ============================================================================
# 测试：守恒归零机制集成
# ============================================================================

class TestConservationToZero:
    """守恒归零机制完整测试"""
    
    def test_conservation_law_integration(self):
        """测试守恒律与公理系统集成"""
        dao = AxiomSystem()
        
        # 定义能量守恒律
        def energy(state: YinYangPair) -> float:
            return abs(state.total())
        
        energy_conservation = ConservationLaw(quantity=energy, name="YinYang Energy")
        dao.A5_conservation.append(energy_conservation)
        
        # 验证系统自洽
        status = dao.verify_all()
        assert status['A5_conservation_ok'], "守恒律验证失败"
    
    def test_ternary_structure_conservation(self):
        """测试三元结构的守恒性"""
        # 定义三元操作
        def creation(x: complex) -> complex:
            return x * (1 + 0.1j)
        
        def annihilation(x: complex) -> complex:
            return x / (1 + 0.1j)
        
        def conservation(x: complex) -> complex:
            return x  # 恒同
        
        ternary = TernaryStructure(
            creation=creation,
            annihilation=annihilation,
            conservation=conservation
        )
        
        # 验证闭合性（守恒回归）
        x = 1.0 + 0.5j
        assert ternary.verify_closure(x), \
            "三元结构不满足守恒闭合性"
    
    def test_multi_system_conservation(self):
        """测试多系统联合守恒"""
        # 创建多个阴阳对
        pairs = [
            YinYangPair(yang=1.0+0j, yin=-1.0+0j),
            YinYangPair(yang=2.0+0j, yin=-2.0+0j),
            YinYangPair(yang=0.5+0j, yin=-0.5+0j),
        ]
        
        # 总守恒量
        def total_conservation(pairs: list[YinYangPair]) -> complex:
            return sum(pair.total() for pair in pairs)
        
        initial_total = total_conservation(pairs)
        
        # 独立演化
        for i in range(len(pairs)):
            for _ in range(50):
                pairs[i] = pairs[i].evolve(rate=0.1)
        
        final_total = total_conservation(pairs)
        
        # 总守恒量不变
        assert abs(initial_total - final_total) < 1e-6, \
            "多系统总守恒量不守恒"
    
    def test_conservation_zero_return_path(self):
        """测试守恒归零路径"""
        pair = YinYangPair(yang=10.0+0j, yin=-10.0+0j)
        
        # 记录演化路径
        path = [pair.total()]
        for _ in range(200):
            pair = pair.evolve(rate=0.05)
            path.append(pair.total())
        
        # 验证路径始终接近零
        max_deviation = max(abs(x) for x in path)
        assert max_deviation < 1e-5, \
            f"演化路径偏离零点过大: {max_deviation}"
        
        # 验证最终归零
        assert pair.to_innate_balance(tolerance=1e-6), \
            "演化未能归零于 ±0"


# ============================================================================
# 性能和边界测试
# ============================================================================

class TestConservationEdgeCases:
    """守恒律边界情况测试"""
    
    def test_zero_initial_state(self):
        """测试零初态的守恒"""
        pair = YinYangPair(yang=0.0+0j, yin=0.0+0j)
        
        for _ in range(100):
            pair = pair.evolve()
        
        assert pair.to_innate_balance(), \
            "零初态应保持归零"
    
    def test_large_amplitude_conservation(self):
        """测试大振幅守恒"""
        pair = YinYangPair(yang=1e6+0j, yin=-1e6+0j)
        initial_total = pair.total()
        
        for _ in range(100):
            pair = pair.evolve()
        
        final_total = pair.total()
        relative_error = abs(final_total - initial_total) / abs(initial_total)
        
        assert relative_error < 1e-6, \
            "大振幅情况守恒精度下降"
    
    def test_complex_phase_conservation(self):
        """测试复相位守恒"""
        pair = YinYangPair(
            yang=1.0 * np.exp(1j * np.pi / 4),
            yin=-1.0 * np.exp(1j * np.pi / 4)
        )
        
        initial_total = pair.total()
        
        for _ in range(100):
            pair = pair.evolve()
        
        final_total = pair.total()
        
        assert abs(final_total - initial_total) < 1e-6, \
            "复相位情况守恒失败"


# ============================================================================
# 运行所有测试
# ============================================================================

if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])
